import os

import pymysql

#获取.env文件中的内容
from dotenv import load_dotenv
load_dotenv()

DATA_BASE_URL=os.getenv("DATA_BASE_URL")
USER_NAME=os.getenv("USER_NAME")
PASSWORD=os.getenv("PASSWORD")
DATABASE_NAME=os.getenv("DATABASE_NAME")


#mysql数据库
class Database:
    def __init__(self, url=DATA_BASE_URL):
        self.url = url
        self.conn = pymysql.connect(host=url, user=USER_NAME, password=PASSWORD, db=DATABASE_NAME,charset='utf8mb4',cursorclass=pymysql.cursors.DictCursor)
        print("数据库已连接")



    def insert_session(self, chat_id, character_id,user_message,system_message,model_name=None):
        """
        创建并保存会话信息
        :param chat_id:会话id
        :param character_id: 角色id
        :param model_name: 使用的大模型名称
        :return:
        """
        with self.conn.cursor() as cursor:
            try:
                sql1 = "INSERT INTO `chat` (`id`, `character_id`, `model_name`) VALUES (%s, %s, %s)"
                cursor.execute(sql1, (chat_id, character_id, model_name))

                #插入system_prompt
                sql2 = "INSERT INTO `chatdetail` (`chat_id`,`message_type` ,`content`) VALUES (%s,%s,%s)"
                cursor.execute(sql2, (chat_id,"system",system_message))

            except Exception as e:
                print(e)

        self.conn.commit()
    def update_session_info(self, chat_id,message_type,content):
        """
        更新会话信息
        :param chat_id:会话id
        :param message_type: 消息类型
        :param content: 消息内容
        :return:
        """
        sql = "INSERT INTO `chatdetail` (`chat_id`,`message_type` ,`content`) VALUES (%s,%s,%s)"
        with self.conn.cursor() as cursor:
            cursor.execute(sql, (chat_id,message_type,content))

        self.conn.commit()
        return "更新成功"

    def select_session_info(self, chat_id):
        """
        查询会话信息 判断会话是否存在
        :param chat_id: 会话id
        :return:
        """

    def select_chat_detail(self, chat_id):
        """
        查询历史记录
        :param chat_id:
        :return:
        """
        sql = "select id,chat_id,message_type,content from chatdetail where chat_id = %s"
        with self.conn.cursor() as cursor:
            cursor.execute(sql, (chat_id))
            result = cursor.fetchall()
            return result

    def select_chat_list(self):
        """
        查询会话列表
        :return:
        """
        sql = "select id,character_id,model_name,create_time from chat ORDER BY create_time DESC"

        with self.conn.cursor() as cursor:
            cursor.execute(sql)
            result = cursor.fetchall()
            return result

    def delete_chat(self,chat_id):
        """
        删除会话
        :param chat_id:
        :return:
        """
        #删除详情表chat信息
        sql = "delete from chatdetail where chat_id = %s"
        with self.conn.cursor() as cursor:
            cursor.execute(sql,(chat_id))


        sql = "delete from chat where id = %s"

        with self.conn.cursor() as cursor:
            cursor.execute(sql,(chat_id))
        self.conn.commit()
        return "删除成功"

    def close(self):
        self.conn.close()
        print("数据库已关闭")


